//
// Created by padeler on 21/6/2017.
//

#ifndef PYOPENPOSE_OPENPOSEWRAPPER_H
#define PYOPENPOSE_OPENPOSEWRAPPER_H

#include <string>
#include <vector>
#include <opencv/cv.h>
#include <memory>
#include <opencv2/opencv.hpp>


class OpenPoseWrapper {


public:
    typedef std::shared_ptr<OpenPoseWrapper> Ptr;
    typedef std::vector<cv::Mat> KeypointGroups;

    enum class ScaleMode : unsigned char
    {
        InputResolution,
        NetOutputResolution,
        OutputResolution,
        ZeroToOne, // [0, 1]
        PlusMinusOne, // [-1, 1]
        UnsignedChar, // [0, 255]
    };

    enum KeypointType { POSE, FACE, HAND };

    /**
     * Openpose wrapper constructor.
     * @param netPoseSize Body pose network input size
     * @param netFaceHandsSize Face and Hands networks input size
     * @param outputSize The output joint positions is rescaled to this size.
     * @param model The model to be used to the body pose estimation
     * @param modelFolder Folder containing the models for body pose, hands and faces (open pose models folder)
     * @param logLevel Openpose loglevel parameter
     * @param downloadHeatmaps If False the heatmaps wont be downloaded from GPU mem.
     * @param heatMapScaleMode The scale mode of the output heatmaps @see ScaleMode
     * @param withFace Initialize the network for face pose detection
     * @param withHands Initialize the network for hand pose detection
     */
    OpenPoseWrapper(const cv::Size &netPoseSize = cv::Size(320, 240), const cv::Size &netFaceHandsSize = cv::Size(128,128),
                    const cv::Size &outputSize = cv::Size(640,480),
                    const std::string &model="COCO", const std::string &modelFolder="models/", const int logLevel=255,
                    bool downloadHeatmaps=false,
                    OpenPoseWrapper::ScaleMode heatMapScaleMode = OpenPoseWrapper::ScaleMode::ZeroToOne,
                    bool withFace=true, bool withHands=true, int gpuId=0);

    /**
     * Detect body poses in given image
     * @param rgb input image
     */
    void detectPose(const cv::Mat &rgb);

    /**
     * Detect faces in given image.
     * You need to call the @refitem detectPose with the same image first.
     *
     * @param rgb input image
     */
    void detectFace(const cv::Mat &rgb);

    /**
     * Detect faces in given image.
     * The network is only run in the supplied rectangles
     * You don't have to call the @refitem detectPose first.
     *
     * @param rgb input image
     * @param faceRects bounding boxes with the hand positions in the image. Nx4 cv::Mat with type CV_32SC1.
     *        One row per person, Each row is a bounding box [x,y,width,height] .
     */
    void detectFace(const cv::Mat &rgb, const cv::Mat &faceRects);

    /**
     * Detect hands in given image.
     * The network is only run in the supplied rectangles
     * You don't have to call the @refitem detectPose first.
     *
     * @param rgb input image
     * @param handRects bounding boxes with the hand positions in the image. Nx8 cv::Mat with type CV_32SC1.
     *        Each row corresponds to  a person. Each row represents two bounding rectangles of the form:
     *         [ leftX, leftY, leftWidth, leftHeight, rightX, rightY, rightWidth, rightHeight]
     */
    void detectHands(const cv::Mat &rgb, const cv::Mat &handRects);


    /**
     * Detect hands in given image.
     * You need to call the @refitem detectPose with the same image first.
     *
     * @param rgb input image
     */
    void detectHands(const cv::Mat &rgb);


    /**
     * Uses the visualization tools of Openpose to render detected keypoints on the input rgb image
     * @param rgb input image
     * @return the input image with the keypoints and skeletons rendered on it.
     */
    cv::Mat render(const cv::Mat &rgb);

    /**
     * Returns the keypoints of a given type.
     * @param t the keypointtype @see KeypointType
     * @return A vector of cv::Mat. The vector has a single entry for t=POSE and t=FACE and two entries for t==HAND.
     * For hands the first entry is a Mat with all left hands, and the second is a Mat with all the right hands.
     * The dimensions of each mat is NxKx3. N is the number of item detected (i.e number of FACEs if t=FACE).
     * K is the number of keypoitns (i.e 21 if t=HAND). The last dimension is the X,Y coordinates of the keypoint and
     * its score.
     */
    KeypointGroups getKeypoints(KeypointType t=POSE);

    /**
     * Returns the heatmaps and PAFs generated by the POSE network.
     * @return a multichannel cv::Mat, each channel corresponds to a different headmap. Each PAF is represented by two
     * consecutive channels one for each dimension of the PAF.
     */
    cv::Mat getHeatmaps();

    /**
     * Returns the hand rectangles (bouding boxes) used during the the
     * last network forward pass (i.e last call to detectHands)
     *
     * The rectangles are stored in a cv::Mat with N rows by 8 columns.
     * Each row corresponds to  a person.
     * Each row represents two bounding rectangles of the form:
     *  [ leftX, leftY, leftWidth, leftHeight, rightX, rightY, rightWidth, rightHeight]
     * if width or height is zero, the corresponding rectangle is ignored.
     * Check the openpose source (handExtractor::forwardPass for more details)
     *
     * @return a cv::Mat Nx8. Each row holds two rectangles, one for the left and one for the right hand.
     */
    const cv::Mat getHandRects() const {
        return handRects;
    }

    /**
     * Returns the face rectangles (bouding boxes) used during the the
     * last network forward pass (i.e last call to detectFace)
     * The rectangles are stored in a cv::Mat with N rows and 4 columns.
     * Each row represents a bounding rectangle of the form:
     *  [ faceX, faceY, faceWidth, faceHeight]
     * Faces with area less than 40 (faceWidth*faceHeight<40) are ignored.
     * Check the openpose source (faceExtractor::forwardPass for more details)
     * @return a vector of rectangles, one for each face.
     */
    const cv::Mat getFaceRects() const {
        return faceRects;
    }


    /**
     * Returns a vector of length 2.
     * Index 0 holds all left hands detected.
     * Index 1 holds all right hands detected.
     * @return all the hands (left and right) of all persons detected during the last call.
     */
    std::vector<cv::Mat> getHandHeatmaps();

    /**
     * Get all face heatmaps for all persons detected.
     * @return a cv::Mat with all heatmaps for all persons.
     */
    cv::Mat getFaceHeatmaps();

private:
    struct PrivateData;
    std::shared_ptr<PrivateData> membersPtr;

    bool withFace, withHands;

    // N rows by 8 columns. Each row corresponds to  a person.
    // Each row represents two bounding rectangles of the form:
    // [ leftX, leftY, leftWidth, leftHeight, rightX, rightY, rightWidth, rightHeight]
    // if width or height is zero, the corresponding rectangle is ignored.
    cv::Mat handRects;

    // N rows by 4 columns. Each row corresponds to  a person.
    // Each row represents a bounding rectangle of the form:
    // [ faceX, faceY, faceWidth, faceHeight]
    cv::Mat faceRects;

};


#endif //PYOPENPOSE_OPENPOSEWRAPPER_H
